package contest.sweek199;

import listbyorder.utils.TreeNode;

import java.util.ArrayList;
import java.util.List;

/**
 * 1530. 好叶子节点对的数量
 * @author code_yc
 * @version 1.0
 * @date 2020/7/26 8:22
 */
public class Solution3 {

    public int countPairs(TreeNode root, int distance) {
        if(root == null || distance < 2) return 0;
        if(root.left == null) return countPairs(root.right, distance);
        if(root.right == null) return countPairs(root.left, distance);
        List<Integer> l = helper(root.left, 1);
        List<Integer> r = helper(root.right, 1);
        int ans = 0;
        for(int n1 : l){
            for(int n2 : r){
                if(n1 + n2 <= distance) ans++;
            }
        }
        return ans += countPairs(root.right, distance) + countPairs(root.left, distance);
    }

    // 所有叶子节点的深度
    private List<Integer> helper(TreeNode root, int depth){
        List<Integer> res = new ArrayList<>();
        if(root.left == null && root.right == null) {
            res.add(depth);
            return res;
        }
        if(root.left != null){
            List<Integer> left = helper(root.left, depth + 1);
            res.addAll(left);
        }
        if(root.right != null){
            List<Integer> right = helper(root.right, depth + 1);
            res.addAll(right);
        }
        return res;
    }
}
